package com.feidee.fd.sml.algorithm.component.feature

import org.apache.spark.ml.PipelineStage
import org.apache.spark.ml.feature.CountVectorizer
import org.apache.spark.sql.DataFrame

/**
  * @Author songhaicheng
  * @Date 2019/1/10 11:17
  * @Description
  * @Reviewer
  */
case class CountVectorizerParam(
                                 override val input_pt: String,
                                 override val output_pt: String,
                                 override val hive_table: String,
                                 override val flow_time: String,
                                 override val inputCol: String,
                                 override val outputCol: String,
                                 override val preserveCols: String,
                                 override val modelPath: String,
                                 // 计算后保存的词数量，> 0（算法默认 2^18^）
                                 vocabSize: Int,
                                 // 词出现在总文本中最小频率数，> 0（算法默认 1.0）
                                 minDF: Double,
                                 // 最小词频，> 0（算法默认 1.0）
                                 minTF: Double,
                                 // 二值化，true 时将所有非 0 的值设为 1，默认 false
                                 binary: Boolean
                               ) extends FeatureParam {

  def this() = this(null, null, null, null, "input", "features", null, null, 1 << 18, 1.0, 1.0, false)

  override def verify(): Unit = {
    super.verify()
    require(vocabSize > 0, "param vocabSize must be grater than 0")
    require(minDF >= 0, "param minDF can't be negative")
    require(minTF >= 0, "param minTF can't be negative")
  }

  override def toMap: Map[String, Any] = {
    var map = super.toMap
    map += ("vocabSize" -> vocabSize)
    map += ("minDF" -> minDF)
    map += ("minTF" -> minTF)
    map += ("binary" -> binary)
    map
  }
}


class CountVectorizerEncoder extends AbstractFeatureEncoder[CountVectorizerParam] {

  override def setUp(param: CountVectorizerParam, data: DataFrame): Array[PipelineStage] = {
    val cnt = new CountVectorizer()
      .setInputCol(param.inputCol)
      .setOutputCol(param.outputCol)
      .setVocabSize(param.vocabSize)
      .setMinDF(param.minDF)
      .setMinTF(param.minTF)
      .setBinary(param.binary)

    Array(cnt)
  }

}

object CountVectorizerEncoder {

  def apply(paramStr: String): Unit = {
    new CountVectorizerEncoder()(paramStr)
  }

  def main(args: Array[String]): Unit = {
    CountVectorizerEncoder(args(0))
  }
}